import numpy as np
import matplotlib.pyplot as plt
from photutils.psf import EPSFModel
from astropy.stats import sigma_clipped_stats
from scipy.optimize import curve_fit

class PSFInspector:
    """
    A class to inspect PSF fitting quality from multiple perspectives:
      - Data / Model / Residual cutouts
      - Residual histogram with Gaussian overlay
      - Radial profiles comparison
      - Reduced chi-square calculation
      - Full-field residual map

    Parameters
    ----------
    data_cutout : 2D ndarray
        The observed image cutout (e.g., around one star).
    model_cutout : 2D ndarray
        The PSF model evaluated at the same size & center.
    sigma_cutout : 2D ndarray
        The per-pixel noise estimate for the cutout.
    full_data : 2D ndarray, optional
        The full observed image (for full-field residuals).
    full_model : 2D ndarray, optional
        The full PSF-model image (same shape as full_data).
    """

    def __init__(self, data_cutout, model_cutout, sigma_cutout,
                 full_data=None, full_model=None):
        self.data = data_cutout
        self.model = model_cutout
        self.sigma = sigma_cutout
        self.residual = self.data - self.model
        self.full_data = full_data
        self.full_model = full_model

    def plot_cutouts(self, vmin_p=5, vmax_p=99, cmap_res='RdBu'):
        """
        Show Data, Model and Residual side by side.
        """
        vmin, vmax = np.percentile(self.data, [vmin_p, vmax_p])
        fig, axes = plt.subplots(1, 3, figsize=(12, 4))
        axes[0].imshow(self.data, origin='lower', vmin=vmin, vmax=vmax)
        axes[0].set_title('Data')
        axes[1].imshow(self.model, origin='lower', vmin=vmin, vmax=vmax)
        axes[1].set_title('PSF Model')
        axes[2].imshow(self.residual, origin='lower',
                       vmin=-vmax*0.2, vmax=vmax*0.2, cmap=cmap_res)
        axes[2].set_title('Residual')
        for ax in axes:
            ax.axis('off')
        plt.tight_layout()
        plt.show()

    def plot_residual_hist(self, bins=50):
        """
        Plot histogram of residuals and overplot a Gaussian fit.
        """
        res = self.residual.flatten()
        res = res[np.isfinite(res)]
        mu, std = np.mean(res), np.std(res)

        plt.figure(figsize=(5,4))
        # histogram
        counts, edges, _ = plt.hist(res, bins=bins, density=True,
                                    alpha=0.6, label='Residuals')
        # Gaussian overlay
        x = np.linspace(edges[0], edges[-1], 300)
        gauss = 1/(std*np.sqrt(2*np.pi)) * np.exp(-(x-mu)**2/(2*std**2))
        plt.plot(x, gauss, '--', label=f'Gaussian fit\nμ={mu:.2f}, σ={std:.2f}')
        plt.xlabel('Residual Flux')
        plt.ylabel('Probability Density')
        plt.legend()
        plt.title('Residual Histogram')
        plt.show()

    def plot_radial_profiles(self, center=None, max_radius=None):
        """
        Compute and plot radial profiles of data vs. model.
        """
        y, x = np.indices(self.data.shape)
        if center is None:
            center = (self.data.shape[0]//2, self.data.shape[1]//2)
        r = np.hypot(x - center[1], y - center[0])
        if max_radius is None:
            max_radius = np.max(r)
        # bin
        nbins = int(max_radius)
        r_bins = np.linspace(0, max_radius, nbins+1)
        digitized = np.digitize(r.flat, r_bins)
        profile_data = np.array([self.data.flat[digitized==i].mean()
                                 for i in range(1, nbins+1)])
        profile_model = np.array([self.model.flat[digitized==i].mean()
                                  for i in range(1, nbins+1)])

        radii = 0.5*(r_bins[:-1] + r_bins[1:])
        plt.figure(figsize=(5,4))
        plt.plot(radii, profile_data/np.max(profile_data),
                 'o', label='Data')
        plt.plot(radii, profile_model/np.max(profile_model),
                 '-', label='Model')
        plt.xlabel('Radius (pix)')
        plt.ylabel('Normalized Flux')
        plt.legend()
        plt.title('Radial Profile')
        plt.show()

    def reduced_chi2(self, nparams=3):
        """
        Compute reduced chi-square:
          χ² = Σ[(D - M)² / σ²]
          χ²_red = χ² / ν, where ν = Npix - nparams
        """
        chi2 = np.sum(((self.residual) / self.sigma)**2)
        ndof = self.data.size - nparams
        chi2_red = chi2 / ndof
        return chi2_red

    def plot_full_residual(self, cmap='RdBu', clip_sigma=3):
        """
        Show the full-field residual map, clipped at +/-clip_sigma*std.
        Requires full_data and full_model to be set.
        """
        if self.full_data is None or self.full_model is None:
            raise ValueError("full_data and full_model must be provided")
        res_full = self.full_data - self.full_model
        std = np.std(res_full)
        vlim = clip_sigma * std
        plt.figure(figsize=(6,6))
        plt.imshow(res_full, origin='lower', cmap=cmap,
                   vmin=-vlim, vmax=vlim)
        plt.colorbar(label='Residual')
        plt.title('Full-field Residuals')
        plt.show()
import numpy as np
import matplotlib.pyplot as plt
from astropy.modeling.models import Gaussian2D
from astropy.stats import sigma_clipped_stats

# 1. 用 Gaussian2D 自己生成一张模拟点源图
size = 101
y, x = np.mgrid[0:size, 0:size]
# 中心在 (50,50)，总振幅 1，sigma=5 pix
gauss = Gaussian2D(amplitude=1, x_mean=50, y_mean=50, x_stddev=5, y_stddev=5)
data = gauss(x, y)

# 加一点背景和噪声
_, median, std = sigma_clipped_stats(data, sigma=3.0)
data += median
noise = np.random.normal(0, std*0.5, size=data.shape)
data += noise



# 3. 构建一个简单的 “模型” 作为示例：用同样的 Gaussian2D，但精确无噪声
model = gauss(x, y)

# 4. 构造噪声估计图（恒定）
sigma_map = np.full_like(data, std)

# 5. 检查
insp = PSFInspector(data, model, sigma_map,
                    full_data=data, full_model=model)
insp.plot_cutouts()
insp.plot_residual_hist()
insp.plot_radial_profiles()
print("Reduced χ² =", insp.reduced_chi2())
insp.plot_full_residual()
